package demo.websocket.server.example3;

import com.alibaba.fastjson.JSONObject;
import org.apache.commons.lang3.StringUtils;
import org.apache.rocketmq.spring.core.RocketMQTemplate;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.MessageHeaders;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.ChannelInterceptorAdapter;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketTransportRegistration;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import java.security.Principal;
import java.util.Map;

@Configuration
@EnableWebSocketMessageBroker
public class StompWebSocketConfig implements WebSocketMessageBrokerConfigurer {

    @Autowired
    private RocketMQTemplate rocketMQTemplate;

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/websocket-sockjs-stomp");

        registry.addEndpoint("/websocket-sockjs-stomp")
//                .addInterceptors(new HttpSessionHandshakeInterceptor())
                .setAllowedOriginPatterns("*")
                .withSockJS();
    }

    @Override
    public void configureWebSocketTransport(WebSocketTransportRegistration registry) {
        registry.setMessageSizeLimit(30 * 1024);
    }

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        registry.enableSimpleBroker("/user", "/topic");
        registry.setApplicationDestinationPrefixes("/app");
        registry.setUserDestinationPrefix("/user");
        registry.setPreservePublishOrder(true);
    }

    /**
     * 1、设置拦截器
     * 2、首次连接的时候，获取其Header信息，利用Header里面的信息进行权限认证
     * 3、通过认证的用户，使用 accessor.setUser(user); 方法，将登陆信息绑定在该 StompHeaderAccessor 上，在Controller方法上可以获取 StompHeaderAccessor 的相关信息
     *
     * @param registration
     */
    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(new ChannelInterceptor() {

            private final Logger logger = LoggerFactory.getLogger(ChannelInterceptorAdapter.class);

            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {
                StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
                //1、判断是否首次连接
                if (StompCommand.CONNECT.equals(accessor.getCommand())) {
                    //2、判断用户名和密码
                    String name = accessor.getNativeHeader("username").get(0);
                    String userId = accessor.getNativeHeader("userId").get(0);

                    if ("admin".equals(name)) {
                        Principal principal = new Principal() {
                            @Override
                            public String getName() {
                                return userId;
                            }
                        };
                        accessor.setUser(principal);
                        accessor.setLogin(userId);
                        Map<String, Object> sessionAttributes = accessor.getSessionAttributes();
                        sessionAttributes.put("userId", userId);
                        return message;
                    } else {
                        return null;
                    }
                } else if (StompCommand.SEND.equals(accessor.getCommand())) {
                    try {
                        String toUserId = accessor.getNativeHeader("toUserId").get(0);
                        Map<String, Object> sessionAttributes = accessor.getSessionAttributes();
                        sessionAttributes.put("toUserId", StringUtils.trimToEmpty(toUserId));
                        sendToRocketMQ(message, accessor);
                    }catch (Exception e){
                        logger.error(e.getMessage(),e);
                    }
                } else if (StompCommand.MESSAGE.equals(accessor.getCommand())) {
                    logger.info("StompCommand.MESSAGE msg={}", message);
                }
                logger.info("------------当前用户={},Command={},msg={}", accessor.getUser().getName(), accessor.getCommand(), message);
                //不是首次连接，已经登陆成功
                return message;
            }

            /**
             * 将来自客户端的消息发送到RocketMQ
             * @param message
             */
            private void sendToRocketMQ(Message<?> message, StompHeaderAccessor accessor) {
                try {
                    String toUserId = (String) accessor.getSessionAttributes().get("toUserId");
                    if (StringUtils.isNotEmpty(toUserId)) {
                        rocketMQTemplate.send("userTopic", message);
                    } else {
                        rocketMQTemplate.send("publicTopic", message);
                    }
                } catch (Exception e) {
                    logger.error(e.getMessage(), e);
                }
            }

            @Override
            public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
//                logger.info("postSend----------{},isSend={}",message,sent);
            }

            @Override
            public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, Exception ex) {
//                logger.info("afterSendCompletion----------ex={}", ex);
            }

            @Override
            public boolean preReceive(MessageChannel channel) {
//                logger.info("preReceive----------");
                return true;
            }

            @Override
            public void afterReceiveCompletion(Message<?> message, MessageChannel channel, Exception ex) {
//                logger.info("afterReceiveCompletion----------");
            }

            @Override
            public Message<?> postReceive(Message<?> message, MessageChannel channel) {
//                logger.info("postReceive msg={}", message);
                return message;
            }

        });
    }
}
